import omp 
import logging
import pprint
import hashlib
from lxml import etree

from config import *

logger = logging.getLogger('dmzwatch.'+__name__)

def openvas_add_default(ip, default_config, default_portlist):
    a = omp.OMPClient()
    a.init_default_config(default_config, default_portlist)
    
    config_postfix = hashlib.md5((default_config + default_portlist).encode('utf-8')).hexdigest()
    portlist_postfix = hashlib.md5(default_portlist.encode('utf-8')).hexdigest()
    
    newtarget = omp.OMPTarget('', TARGET_PREFIX + '_' + ip + '_' + portlist_postfix)
    newtask = omp.OMPTask('', TASK_PREFIX + '_' + ip + '_' + config_postfix)
    
    target_found = False
    logger.debug('Target: %s', newtarget)
    logger.debug('Task: %s', newtask)
    for target in a.get_targets():
        if target.name == newtarget.name:
            logger.debug('Target is here')
            target_found = True
            newtarget.uid = target.uid
            break
            
    if target_found == False:
        logger.debug('Adding new target')
        newtarget.uid = a.create_target('127.0.0.1', newtarget.name, a.defaultportlist.uid)
    
    task_found = False
    for task in a.get_tasks():    
        if task.name == newtask.name:
            logger.debug('Task is here')
            task_found = True
            newtask.uid = task.uid
            break
            
    if task_found == False:
        logger.debug('Adding new task')
        newtask.uid = a.create_task(newtask.name, newtarget.uid, a.defaultconfig.uid)
        
    status = a.get_status(newtask.uid)
    logger.debug('Status is %s', status)
    
    if status == omp.TaskStatus.running or status == omp.TaskStatus.requested:
        raise Exception('Task is running already')
    else:
        logger.debug('Running new task')
        report_uid = a.run_task(newtask.uid)
        
    return {'task' : newtask.uid, 'report': report_uid}

def openvas_tasks_number():
    a = omp.OMPClient()
    num = 0
    for task in a.get_tasks():
        if task.status == omp.TaskStatus.requested or task.status == omp.TaskStatus.running:
            num = num + 1
    return num
    
def openvas_check_task(task_uid):    
    a = omp.OMPClient()
    status = a.get_status(task_uid)
    
    if status == omp.TaskStatus.new:
        a.run_task(newtask.uid)
        return ExecStatus.running
        
    elif status == omp.TaskStatus.requested:
        return ExecStatus.running
        
    elif status == omp.TaskStatus.running:
        return ExecStatus.running    
        
    elif status == omp.TaskStatus.done:
        return ExecStatus.complete
        
    elif status == omp.TaskStatus.error:
        return ExecStatus.error

    else:
         return ExecStatus.error
 

def openvas_get_reports(report_uid):
    a = omp.OMPClient()
    a.init_default_reports()
    return a.get_reports(report_uid)

needed_keys = ['host', 'port', 'threat', 'nvt/name',  'nvt/family', 'nvt/cvss_base', 'nvt/risk_factor', 'nvt/cve', 'nvt/bid', 'nvt/cert', 'nvt/tags', 'nvt/xref', 'description' ] 
ThreatLevel = enum (
            log = 'Log',
            low = 'Low',
            medium = 'Medium',
            high = 'High' )
ThreatNum = {
                ThreatLevel.log         : '0',
                ThreatLevel.low         : '1',
                ThreatLevel.medium      : '2',
                ThreatLevel.high        : '3'
            }       
def parse_xml(xml): 
    root = etree.fromstring(xml.decode('utf-8'))
    vulns = []
    for result in root.findall('report/results/result'):
        vulnerability = {}
        if result.find('threat').text == ThreatLevel.log:
            continue
        for key in needed_keys:
            try:
                value = result.find(key).text
            except:
                vulnerability.update({key.split('/')[-1] : 'None'})
            else:
                vulnerability.update({key.split('/')[-1] : value})
        vulnerability.update({'threat': ThreatNum[vulnerability['threat']]})        
        vulns.append(vulnerability)
    return vulns
    
    
    
if __name__ == '__main__':
    with open('openvas.xml', 'rb') as f:
        xml = f.read()
    pprint.pprint(parse_xml(xml))
    
    
